package com.kongjs.common.security.component;

import org.springframework.security.access.hierarchicalroles.RoleHierarchy;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import java.util.Collection;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/*
 * 使用需要实现角色或者权限继承
 * https://docs.spring.io/spring-security/reference/servlet/authorization/method-security.html
 * */
@Component
public class DynamicRoleHierarchy implements RoleHierarchy {

    private final Map<String, Set<GrantedAuthority>> rolesReachableInOneOrMoreStepsMap = new ConcurrentHashMap<>();
    private final Map<String, Set<GrantedAuthority>> rolePermissionsReachableInOneOrMoreStepsMap = new ConcurrentHashMap<>();

    public static String removePrefix(String str, String prefix) {
        if (str != null && str.startsWith(prefix)) {
            return str.substring(prefix.length());
        }
        return str;
    }

    @Override
    public Collection<? extends GrantedAuthority> getReachableGrantedAuthorities(Collection<? extends GrantedAuthority> authorities) {
        Set<GrantedAuthority> grantedAuthoritiesAll = new HashSet<>();
        if (CollectionUtils.isEmpty(authorities)) {
            return grantedAuthoritiesAll;
        }
        for (GrantedAuthority role : authorities) {
            if (CollectionUtils.isEmpty(rolesReachableInOneOrMoreStepsMap)) {
                break;
            }
            Set<GrantedAuthority> roleAuthorities = rolesReachableInOneOrMoreStepsMap.get(removePrefix(role.getAuthority(), "SCOPE_"));
            if (CollectionUtils.isEmpty(roleAuthorities)) {
                continue;
            }
            grantedAuthoritiesAll.addAll(roleAuthorities);
            for (GrantedAuthority permission : roleAuthorities) {
                if (CollectionUtils.isEmpty(rolePermissionsReachableInOneOrMoreStepsMap)) {
                    break;
                }
                Set<GrantedAuthority> permissionAuthorities = rolePermissionsReachableInOneOrMoreStepsMap.get(removePrefix(permission.getAuthority(), "SCOPE_"));
                if (CollectionUtils.isEmpty(permissionAuthorities)) {
                    continue;
                }
                grantedAuthoritiesAll.addAll(permissionAuthorities);
            }
        }
        grantedAuthoritiesAll.addAll(authorities);
        grantedAuthoritiesAll = grantedAuthoritiesAll.stream().map(grantedAuthority -> new SimpleGrantedAuthority(removePrefix(grantedAuthority.getAuthority(), "SCOPE_"))).collect(Collectors.toSet());
        return grantedAuthoritiesAll;
    }

    public void setRolesReachableInOneOrMoreStepsMap(Map<String, Set<GrantedAuthority>> rolesReachableInOneOrMoreStepsMap) {
        this.rolesReachableInOneOrMoreStepsMap.clear();
        this.rolesReachableInOneOrMoreStepsMap.putAll(rolesReachableInOneOrMoreStepsMap);
    }
}
